import os
import torch
import torchvision
from torch import nn
from torch import Tensor
from PIL import Image

RGB_STD = torch.Tensor([0.229, 0.224, 0.225])
RGB_MEAN = torch.Tensor([0.485, 0.456, 0.406])


def preprocess_image(image, shape) -> Tensor:
    """Preprocesses a image to a tensor of a batch size input.

    In Short the images are converted from RGB to BGR,
    then each color channel is zero-centered with respect to the ImageNet dataset,
    without scaling and mean values used for zero-centering each channel are [103.939, 116.779, 123.68].

    Args:
        image (PIL Image or Tensor): Image to be converted to tensor.
        shape (sequence or int): Desired output size.

    .. _references: https://askdevz.com/question/514538-keras-style-transfer-remove-zero-center-by-mean-pixel,
                    https://github.com/keras-team/keras-applications/blob/master/keras_applications/imagenet_utils.py
    """
    transforms = torchvision.transforms.Compose([
        torchvision.transforms.Resize(shape),
        torchvision.transforms.ToTensor(),
        torchvision.transforms.Normalize(mean=RGB_MEAN, std=RGB_STD)])
    return transforms(image).unsqueeze(0)  # (batch_size(1), channel, height, width)


def postprocess_image(image) -> Image:
    """

    In this function, the same mean values ([103.939, 116.779, 123.68]) are added to each respective channel and
    then converted back to the original form, from 'BGR' -> 'RGB'.
    """
    image = image[0].to(RGB_STD.device)
    image = torch.clamp(image.permute(1, 2, 0) * RGB_STD + RGB_MEAN, 0, 1)
    return torchvision.transforms.ToPILImage()(image.permute(2, 0, 1))


def get_pretrained_net(net_name, content_layers, style_layers, device=None):
    """
    """
    supported_net = ("vgg19",)
    assert net_name in supported_net, f"The supported nets are {', '.join(supported_net)}."

    if net_name == "vgg19":
        pretrained_net = torchvision.models.vgg19(pretrained=True)
        return nn.Sequential(*[pretrained_net.features[i]
                               for i in range(max(content_layers + style_layers) + 1)]).to(device)
    return None


def get_gpu(device_count=0):
    """Return gpu(i) if exists, otherwise return cpu().
    """
    if torch.cuda.device_count() >= device_count + 1:
        return torch.device(f"cuda:{device_count}")
    return torch.device("cpu")


def make_dir(directory, recursive=False):
    """Folder creation, ignoring existing folders.
    """
    if not os.path.exists(directory):
        if recursive:
            os.mkdir(directory)
        else:
            os.makedirs(directory)


def get_output_config(work_name=None):
    """Style picture training save picture location.
    """
    make_dir("./checkpoints")
    if work_name is not None:
        output_dir = f"./checkpoints/{work_name}"
        make_dir(output_dir, True)
        return work_name, output_dir

    last_name = sorted(filter(lambda x: x.startswith("nst"), os.listdir("./checkpoints")), reverse=True)
    if last_name is None or len(last_name) == 0:
        work_name = "nst-01"
        output_dir = f"./checkpoints/{work_name}"
    else:
        work_name = "nst-{:0>2d}".format(int(last_name[0].split('-')[-1]) + 1)
        output_dir = f"./checkpoints/{work_name}"

    make_dir(output_dir)
    return work_name, output_dir
